#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
YOLO11 检测模型，实现推理与后处理逻辑。
移植自 legacy src/detection_core/yolo_detect.py，并适配新的配置系统。
"""

from time import time
from typing import Iterable, Tuple

import cv2
import numpy as np
from scipy.special import softmax
import logging

from .base_model import BaseModel
from ...config.settings import ModelConfig, DetectionConfig

logger = logging.getLogger("RDK_YOLO")


class YOLO11Detect(BaseModel):
    """
    YOLO11 检测模型类，负责模型推理输出的解码、后处理与坐标还原。
    """

    def __init__(
        self,
        model_file: str,
        model_config: ModelConfig,
        conf: float,
        iou: float,
    ):
        super().__init__(model_file)

        self.model_config = model_config
        self.input_image_size = model_config.input_image_size
        self.conf = conf
        self.iou = iou
        self.conf_inverse = -np.log(1 / conf - 1)
        logger.info("iou threshold = %.2f, conf threshold = %.2f", iou, conf)
        logger.info("sigmoid_inverse threshold = %.2f", self.conf_inverse)

        # 反量化系数
        self.s_bboxes_scale = self.quantize_model[0].outputs[0].properties.scale_data[np.newaxis, :]
        self.m_bboxes_scale = self.quantize_model[0].outputs[1].properties.scale_data[np.newaxis, :]
        self.l_bboxes_scale = self.quantize_model[0].outputs[2].properties.scale_data[np.newaxis, :]
        logger.info("%s, %s, %s", self.s_bboxes_scale.shape, self.m_bboxes_scale.shape, self.l_bboxes_scale.shape)

        # DFL 求期望的系数
        self.dfl_bins = model_config.dfl_bins
        self.weights_static = np.arange(self.dfl_bins, dtype=np.float32)[np.newaxis, np.newaxis, :]
        logger.info("DFL bins weights shape: %s", self.weights_static.shape)

        # anchor 初始化
        anchor_shapes = model_config.anchor_shapes or [
            {"size": 40, "stride": 8},
            {"size": 20, "stride": 16},
            {"size": 10, "stride": 32},
        ]
        self.s_anchor = self._build_anchor(anchor_shapes[0]["size"])
        self.m_anchor = self._build_anchor(anchor_shapes[1]["size"])
        self.l_anchor = self._build_anchor(anchor_shapes[2]["size"])
        logger.info("anchor shapes: %s %s %s", self.s_anchor.shape, self.m_anchor.shape, self.l_anchor.shape)

    @staticmethod
    def _build_anchor(size: int) -> np.ndarray:
        """构建 anchor 网格。"""
        return np.stack(
            [
                np.tile(np.linspace(0.5, size - 0.5, size), reps=size),
                np.repeat(np.arange(0.5, size + 0.5, 1), size),
            ],
            axis=0,
        ).transpose(1, 0)

    def post_process(self, outputs: Iterable[np.ndarray]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        对模型输出进行后处理，包括解码、阈值筛选、NMS、坐标还原等。

        Returns:
            (class_ids, scores, bboxes)
        """
        begin_time = time()
        outputs = list(outputs)
        # reshape 修改为对应的尺寸
        s_bboxes = outputs[1].reshape(-1, 64)  # 40x40 -> 1600
        s_clses = outputs[0].reshape(-1, 1)

        m_bboxes = outputs[3].reshape(-1, 64)  # 20x20 -> 400
        m_clses = outputs[2].reshape(-1, 1)

        l_bboxes = outputs[5].reshape(-1, 64)  # 10x10 -> 100
        l_clses = outputs[4].reshape(-1, 1)

        # 计算每个尺度的最大置信度及有效索引
        s_max_scores = np.max(s_clses, axis=1)
        s_valid_indices = np.flatnonzero(s_max_scores >= self.conf_inverse)
        s_valid_indices = s_valid_indices[: s_bboxes.shape[0]]

        m_max_scores = np.max(m_clses, axis=1)
        m_valid_indices = np.flatnonzero(m_max_scores >= self.conf_inverse)
        m_valid_indices = m_valid_indices[: m_bboxes.shape[0]]

        l_max_scores = np.max(l_clses, axis=1)
        l_valid_indices = np.flatnonzero(l_max_scores >= self.conf_inverse)
        l_valid_indices = l_valid_indices[: l_bboxes.shape[0]]

        # 获取每个有效目标的类别 ID 和置信度
        s_ids = np.argmax(s_clses[s_valid_indices, :], axis=1)
        s_scores = s_max_scores[s_valid_indices]

        m_ids = np.argmax(m_clses[m_valid_indices, :], axis=1)
        m_scores = m_max_scores[m_valid_indices]

        l_ids = np.argmax(l_clses[l_valid_indices, :], axis=1)
        l_scores = l_max_scores[l_valid_indices]

        # Sigmoid
        s_scores = 1 / (1 + np.exp(-s_scores))
        m_scores = 1 / (1 + np.exp(-m_scores))
        l_scores = 1 / (1 + np.exp(-l_scores))

        # Bounding Box 分支
        s_bboxes_float32 = s_bboxes[s_valid_indices, :]
        m_bboxes_float32 = m_bboxes[m_valid_indices, :]
        l_bboxes_float32 = l_bboxes[l_valid_indices, :]

        # dist2bbox (ltrb2xyxy)
        s_ltrb_indices = np.sum(
            softmax(s_bboxes_float32.reshape(-1, 4, self.dfl_bins), axis=2) * self.weights_static,
            axis=2,
        )
        s_anchor_indices = self.s_anchor[s_valid_indices, :]
        s_x1y1 = s_anchor_indices - s_ltrb_indices[:, 0:2]
        s_x2y2 = s_anchor_indices + s_ltrb_indices[:, 2:4]
        s_dbboxes = np.hstack([s_x1y1, s_x2y2]) * 8

        m_ltrb_indices = np.sum(
            softmax(m_bboxes_float32.reshape(-1, 4, self.dfl_bins), axis=2) * self.weights_static,
            axis=2,
        )
        m_anchor_indices = self.m_anchor[m_valid_indices, :]
        m_x1y1 = m_anchor_indices - m_ltrb_indices[:, 0:2]
        m_x2y2 = m_anchor_indices + m_ltrb_indices[:, 2:4]
        m_dbboxes = np.hstack([m_x1y1, m_x2y2]) * 16

        l_ltrb_indices = np.sum(
            softmax(l_bboxes_float32.reshape(-1, 4, self.dfl_bins), axis=2) * self.weights_static,
            axis=2,
        )
        l_anchor_indices = self.l_anchor[l_valid_indices, :]
        l_x1y1 = l_anchor_indices - l_ltrb_indices[:, 0:2]
        l_x2y2 = l_anchor_indices + l_ltrb_indices[:, 2:4]
        l_dbboxes = np.hstack([l_x1y1, l_x2y2]) * 32

        # 将三个尺度拼接
        dbboxes = np.concatenate((s_dbboxes, m_dbboxes, l_dbboxes), axis=0)
        scores = np.concatenate((s_scores, m_scores, l_scores), axis=0)
        ids = np.concatenate((s_ids, m_ids, l_ids), axis=0)

        # NMS
        indices = cv2.dnn.NMSBoxes(dbboxes.tolist(), scores.tolist(), self.conf, self.iou)

        if len(indices) == 0:
            return np.array([]), np.array([]), np.array([])

        indices = indices.flatten()
        # 还原到原始图像尺度
        bboxes = dbboxes[indices] * np.array([self.x_scale, self.y_scale, self.x_scale, self.y_scale])
        bboxes = bboxes.astype(np.int32)

        # logger.debug("\033[1;31m" + f"Post Process time = {1000*(time() - begin_time):.2f} ms" + "\033[0m")
        return ids[indices], scores[indices], bboxes


def draw_detection(
    img: np.ndarray,
    bbox: Tuple[int, int, int, int],
    score: float,
    class_id: int,
    detection_config: DetectionConfig,
    tracking_enabled: bool = False,
) -> None:
    """
    在图像上绘制检测框和标签。
    """

    x1, y1, x2, y2 = bbox
    class_names = detection_config.class_names or ["face"]

    if tracking_enabled:
        color = (0, 165, 255)  # 橙色
        label = f"{class_names[class_id]}: {score:.2f} (Tracking)"
    else:
        color = (0, 255, 0)  # 绿色
        label = f"{class_names[class_id]}: {score:.2f}"

    cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
    (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    label_x, label_y = x1, y1 - 10 if y1 - 10 > label_height else y1 + 10
    cv2.rectangle(
        img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
    )
    cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
